import os
import torch
import torch.nn.functional as F
from synthesizers.synthesizer import Synthesizer
from tasks.task import Task


class WaNetParameters:
    def __init__(self, dataset='imagenet10', input_channel=3, input_height=224, input_width=224,
                 attack_mode='all2one', s=0.5, k=4, grid_rescale=1.0, device='cuda'):
        self.dataset = dataset
        self.input_channel = input_channel
        self.input_height = input_height
        self.input_width = input_width
        self.attack_mode = attack_mode
        self.s = s
        self.k = k
        self.grid_rescale = grid_rescale
        self.device = device

        # paths
        self.checkpoints = './synthesizers/triggers/wanet_checkpoints'
        self.ckpt_folder = os.path.join(self.checkpoints, self.dataset)
        self.ckpt_path = os.path.join(
            self.ckpt_folder,
            f"{self.dataset}_{self.attack_mode}_morph.pth.tar"
        )
        self.log_dir = os.path.join(self.ckpt_folder, 'log_dir')


class WaNetSynthesizer(Synthesizer):
    def __init__(self, task: Task, dataset: str = 'imagenet10'):
        super().__init__(task)
        # Expect task.params.input_shape = (3, 224, 224)
        self.task = task
        self.opt = WaNetParameters(
            dataset=dataset,
            input_channel=task.params.input_shape[0],
            input_height=task.params.input_shape[1],
            input_width=task.params.input_shape[2],
            device=task.params.device
        )
        self.dataset = dataset

        # load precomputed grids
        self.identity_grid, self.noise_grid = self._load_grids()

    def _load_grids(self):
        if not os.path.exists(self.opt.ckpt_path):
            raise FileNotFoundError(f"Checkpoint not found: {self.opt.ckpt_path}")
        state = torch.load(self.opt.ckpt_path, map_location=self.opt.device)
        identity = state['identity_grid']
        noise = state['noise_grid']
        return identity, noise

    def synthesize_inputs(self, batch, attack_portion=None):
        """
        Apply the WaNet backdoor to the first `attack_portion` samples in the batch.
        """
        # prepare base grid
        base = (self.identity_grid + self.opt.s * self.noise_grid / self.opt.input_height) * self.opt.grid_rescale
        base = torch.clamp(base, -1, 1)

        # random perturbation
        rand = torch.rand(batch.batch_size,
                          self.opt.input_height,
                          self.opt.input_width, 2,
                          device=self.opt.device) * 2 - 1
        perturbed = torch.clamp(base.unsqueeze(0).repeat(batch.batch_size, 1, 1, 1) + rand / self.opt.input_height,
                                -1, 1)

        # apply grid_sample on attack subset
        if attack_portion is None:
            attack_portion = batch.batch_size
        batch.inputs[:attack_portion] = F.grid_sample(
            batch.inputs[:attack_portion],
            base.unsqueeze(0).repeat(attack_portion, 1, 1, 1),
            align_corners=True
        )
        return batch

    def synthesize_labels(self, batch, attack_portion=None):
        # keep original labels or redirect for all2one
        return batch.labels

    def apply_backdoor_to_a_sample(self, data, label, params):
        data = data.unsqueeze(0)  # 添加 batch 维度，变为 (1, channels, height, width)
        """将backdoor模式应用到样本中"""
        grid_temps = (self.identity_grid + self.opt.s * self.noise_grid / self.opt.input_height) * self.opt.grid_rescale
        grid_temps = torch.clamp(grid_temps, -1, 1)

        # ins = torch.rand(1, self.opt.input_height, self.opt.input_height, 2).to(self.opt.device) * 2 - 1
        ins = torch.rand(1, self.opt.input_height, self.opt.input_width, 2, device=self.opt.device)*2 - 1

        grid_temps2 = grid_temps.repeat(1, 1, 1, 1) + ins / self.opt.input_height
        grid_temps2 = torch.clamp(grid_temps2, -1, 1)
        
        backdoor_sample = F.grid_sample(data, grid_temps.repeat(1, 1, 1, 1), align_corners=True)
        backdoor_sample = backdoor_sample.squeeze(0)  # 去掉 batch 维度
        
        return backdoor_sample



    # def apply_backdoor_to_a_sample(self, data, label=None, params=None):
    #     """
    #     Apply the WaNet backdoor to a single sample tensor.
    #     Input: data of shape (C, H, W)
    #     Returns: backdoored sample of same shape
    #     """
    #     data = data.unsqueeze(0).to(self.opt.device)
    #     base = (self.identity_grid + self.opt.s * self.noise_grid / self.opt.input_height) * self.opt.grid_rescale
    #     base = torch.clamp(base, -1, 1)

    #     backdoored = F.grid_sample(data, base.unsqueeze(0), align_corners=True)
    #     return backdoored.squeeze(0)

    # def apply_backdoor_to_a_sample(self, data, label=None, params=None):
    #     """
    #     Apply the WaNet backdoor to a single sample tensor.
    #     Input: data of shape (C, H, W)
    #     Returns: backdoored sample of same shape
    #     """
    #     # 1) 扩 batch 维度
    #     data = data.unsqueeze(0).to(self.opt.device)  # now (1, C, H, W)

    #     # 2) 计算 base
    #     base = (self.identity_grid + self.opt.s * self.noise_grid / self.opt.input_height) \
    #         * self.opt.grid_rescale
    #     base = torch.clamp(base, -1, 1)
    #     # 假设 base 原本可能是 (H, W, 2) 或 (1, H, W, 2)

    #     # 3) 统一到 (1, H, W, 2)
    #     if base.dim() == 3:
    #         grid = base.unsqueeze(0)
    #     elif base.dim() == 4:
    #         # 如果是 (1, H, W, 2)，直接用它
    #         grid = base
    #     else:
    #         raise ValueError(f"Unexpected grid shape: {base.shape}")

    #     # 4) 调用 grid_sample
    #     backdoored = F.grid_sample(
    #         data,       # (1, C, H, W)
    #         grid,       # now (1, H, W, 2)
    #         align_corners=True
    #     )

    #     # 5) squeeze 拿回 (C, H, W)
    #     return backdoored.squeeze(0)
